from typing import Union, Any
import torch
from feasytools import ArgChecker

from .Pyraformer import Model as Pyraformer
from .DNN_LR import Model as DNN
from .LSTM_LR import Model as LSTM
from .Transformer_LR import Model as Transformer
from .wargs import WorkerArguments

def get_model_save_dir(opt:WorkerArguments):
    model_name = opt.model
    if opt.model != "DNN": model_name += "_" + opt.decoder
    model_save_dir = 'ckpts/{}/{}/i{}_o{}_S{}_f{}'.format(opt.data, model_name,
        opt.ilen, opt.olen, opt.sr, 0 if opt.non_future else 1)
    if opt.model.endswith("former"):
        model_save_dir += f"_nh{opt.n_head}"
    if opt.non_future:
        model_save_dir += "_nf"
    model_save_dir += "/"
    return model_save_dir

def create_model(opt:WorkerArguments) -> torch.nn.Module:
    model_list = {
        "DNN": DNN,
        "LSTM": LSTM,
        "Transformer": Transformer,
        "Pyraformer": Pyraformer,
    }
    model:torch.nn.Module = model_list[opt.model](opt) 
    model.to(opt.device)
    return model
